package com.chatmcp.mcprouter.controller;

import com.chatmcp.mcprouter.dto.ProxyInfoDto;
import com.chatmcp.mcprouter.dto.ServerConfigDto;
import com.chatmcp.mcprouter.jsonrpc.JsonRpcError;
import com.chatmcp.mcprouter.jsonrpc.JsonRpcRequest;
import com.chatmcp.mcprouter.jsonrpc.JsonRpcResponse;
import com.chatmcp.mcprouter.jsonrpc.InitializeParams;
import com.chatmcp.mcprouter.jsonrpc.InitializeResult;
import com.chatmcp.mcprouter.service.McpClientService;
import com.chatmcp.mcprouter.service.McpServerService;
import com.chatmcp.mcprouter.service.ProxyService;
import com.chatmcp.mcprouter.service.impl.ProxyServiceImpl;
import com.chatmcp.mcprouter.util.HashUtil;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import javax.servlet.http.HttpServletRequest;
import java.io.IOException;
import java.time.LocalDateTime;
import java.util.Map;
import java.util.HashMap;
import java.util.UUID;

import com.chatmcp.mcprouter.service.SessionManager;
import com.chatmcp.mcprouter.model.SseSession;

/**
 * 代理控制器
 * 处理MCP代理请求和SSE连接
 */
@Slf4j
@RestController
@RequiredArgsConstructor
public class ProxyController {

    private final ProxyService proxyService;
    private final McpServerService mcpServerService;
    private final McpClientService mcpClientService;
    private final ObjectMapper objectMapper;
    private final SessionManager sessionManager;

    /**
     * SSE连接端点
     *
     * @param key     服务器键
     * @param request HTTP请求
     * @return SSE发射器
     */
    @GetMapping(value = "/sse/{key}", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter sse(@PathVariable String key, HttpServletRequest request) {
        // 检查服务器键是否有效
        ServerConfigDto serverConfig = mcpServerService.getServerConfig(key);
        if (serverConfig == null) {
            throw new IllegalArgumentException("Invalid server key: " + key);
        }

        // 创建会话ID - 使用与Go实现一致的方式
        String sessionId = HashUtil.md5(key + "_" + UUID.randomUUID().toString());
        
        // 创建代理信息
        ProxyInfoDto proxyInfo = ProxyInfoDto.builder()
                .sessionId(sessionId)
                .requestFrom(request.getRemoteAddr())
                .serverKey(key)
                .serverUuid(serverConfig.getServerUuid())
                .serverName(serverConfig.getServerName())
                .serverType(serverConfig.getServerType())
                .serverUrl(serverConfig.getServerUrl())
                .serverCommand(serverConfig.getCommand())
                .serverCommandHash(serverConfig.getCommandHash())
                .serverShareProcess(serverConfig.getShareProcess())
                .connectionTime(LocalDateTime.now())
                .build();
        
        // 使用SessionManager创建和管理会话
        SseSession session = sessionManager.createSession(key, serverConfig, proxyInfo);
        
        // 启动心跳调度器 - 每30秒发送一次心跳
        sessionManager.startHeartbeat(sessionId, 30000);

        // 发送会话ID
        try {
            Map<String, String> data = new HashMap<>();
            data.put("sessionId", sessionId);
            sessionManager.sendEvent(sessionId, "session", data);
        } catch (Exception e) {
            log.error("Failed to send session ID: {}", e.getMessage());
            sessionManager.closeSession(sessionId);
            throw new RuntimeException("Failed to initialize session", e);
        }

        return session.getEmitter();
    }

    /**
     * 处理SSE消息
     *
     * @param request HTTP请求
     * @param body    请求体
     * @return 响应
     */
    @PostMapping("/messages")
    public ResponseEntity<String> messages(HttpServletRequest request, @RequestBody String body) {
        try {
            // 解析请求
            Map<String, Object> requestMap = objectMapper.readValue(body, Map.class);
            String sessionId = (String) requestMap.get("sessionId");
            String message = (String) requestMap.get("message");
            
            if (sessionId == null || message == null) {
                return ResponseEntity.badRequest().body("Missing sessionId or message");
            }
            
            // 获取会话
            SseSession session = sessionManager.getSession(sessionId);
            if (session == null) {
                log.warn("No session found for ID: {}", sessionId);
                return ResponseEntity.notFound().build();
            }
            
            // 获取代理信息
            ProxyInfoDto proxyInfo = session.getProxyInfo();
            if (proxyInfo == null) {
                log.warn("No proxy info found for session: {}", sessionId);
                return ResponseEntity.badRequest().body("Invalid session");
            }
            
            // 处理消息并发送
            String response = proxyService.processSseMessage(proxyInfo.getServerKey(), message, proxyInfo);
            sessionManager.sendEvent(sessionId, "message", response);
            
            return ResponseEntity.ok("Message sent");
        } catch (Exception e) {
            log.error("Failed to process message: {}", e.getMessage());
            return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR)
                    .body("Error: " + e.getMessage());
        }
    }

    /**
     * 处理MCP请求
     *
     * @param key     服务器键
     * @param request HTTP请求
     * @param body    请求体
     * @return 响应
     */
    @RequestMapping(value = "/mcp/{key}", method = {RequestMethod.GET, RequestMethod.POST, RequestMethod.PUT, RequestMethod.DELETE})
    public ResponseEntity<Object> mcp(@PathVariable String key, HttpServletRequest request, @RequestBody(required = false) String body) {
        try {
            // 检查服务器键是否有效
            ServerConfigDto serverConfig = mcpServerService.getServerConfig(key);
            if (serverConfig == null) {
                return ResponseEntity.badRequest().body(
                        JsonRpcResponse.invalidRequest("Invalid server key"));
            }
            
            // 获取会话ID
            String sessionId = request.getHeader("Mcp-Session-Id");
            
            // 解析JSON-RPC请求
            JsonRpcRequest jsonRpcRequest = null;
            if (StringUtils.isNotBlank(body)) {
                try {
                    jsonRpcRequest = objectMapper.readValue(body, JsonRpcRequest.class);
                } catch (JsonProcessingException e) {
                    log.error("Failed to parse JSON-RPC request: {}", e.getMessage());
                    return ResponseEntity.badRequest().body(
                            JsonRpcResponse.parseError(null));
                }
            }
            
            // 检查是否为通知
            if (jsonRpcRequest != null && (jsonRpcRequest.getResult() != null || jsonRpcRequest.getError() != null)) {
                // 这是一个通知而不是请求，直接返回成功
                return ResponseEntity.ok().build();
            }
            
            // 创建代理信息对象
            ProxyInfoDto proxyInfo = ProxyInfoDto.builder()
                    .sessionId(sessionId)
                    .requestId(jsonRpcRequest != null ? jsonRpcRequest.getId() : null)
                    .requestFrom(request.getRemoteAddr())
                    .serverKey(key)
                    .serverUuid(serverConfig.getServerUuid())
                    .serverName(serverConfig.getServerName())
                    .serverType(serverConfig.getServerType())
                    .serverUrl(serverConfig.getServerUrl())
                    .serverCommand(serverConfig.getCommand())
                    .serverCommandHash(serverConfig.getCommandHash())
                    .serverShareProcess(serverConfig.getShareProcess())
                    .jsonrpcVersion(jsonRpcRequest != null ? jsonRpcRequest.getJsonrpc() : null)
                    .requestMethod(jsonRpcRequest != null ? jsonRpcRequest.getMethod() : null)
                    .requestParams(jsonRpcRequest != null ? jsonRpcRequest.getParams() : null)
                    .requestTime(LocalDateTime.now())
                    .build();
            
            // 如果是初始化请求，处理会话ID
            if (jsonRpcRequest != null && "initialize".equals(jsonRpcRequest.getMethod())) {
                try {
                    // 解析初始化参数
                    InitializeParams params = objectMapper.convertValue(jsonRpcRequest.getParams(), InitializeParams.class);
                    
                    // 生成新的会话ID - 使用MD5哈希与Golang实现一致
                    sessionId = HashUtil.md5(key);
                    
                    // 更新代理信息
                    proxyInfo.setSessionId(sessionId);
                    proxyInfo.setConnectionTime(LocalDateTime.now());
                    proxyInfo.setClientName(params.getClientInfo().getName());
                    proxyInfo.setClientVersion(params.getClientInfo().getVersion());
                    proxyInfo.setProtocolVersion(params.getProtocolVersion());
                    
                    // 如果存在对应的SSE会话，更新其信息
                    SseSession session = sessionManager.getSession(sessionId);
                    if (session != null) {
                        session.setProxyInfo(proxyInfo);
                    }
                    
                    // 存储代理信息
                    proxyService.storeProxyInfo(sessionId, proxyInfo);
                } catch (Exception e) {
                    log.error("Failed to process initialize request: {}", e.getMessage(), e);
                    return ResponseEntity.badRequest().body(
                            JsonRpcResponse.parseError(jsonRpcRequest.getId()));
                }
            } else if (StringUtils.isBlank(sessionId)) {
                // 非初始化请求但没有会话ID
                return ResponseEntity.badRequest().body(
                        JsonRpcResponse.invalidRequest("Missing Mcp-Session-Id header"));
            } else {
                // 使用SessionManager获取现有会话
                SseSession session = sessionManager.getSession(sessionId);
                
                // 加载已有的会话信息
                ProxyInfoDto existingInfo = session != null ? session.getProxyInfo() : proxyService.getProxyInfo(sessionId);
                if (existingInfo != null) {
                    // 更新现有信息
                    proxyInfo.setClientName(existingInfo.getClientName());
                    proxyInfo.setClientVersion(existingInfo.getClientVersion());
                    proxyInfo.setProtocolVersion(existingInfo.getProtocolVersion());
                    proxyInfo.setConnectionTime(existingInfo.getConnectionTime());
                    
                    // 如果存在会话，更新其代理信息
                    if (session != null) {
                        session.setProxyInfo(proxyInfo);
                    }
                }
            }
            
            // 连接到MCP服务器并转发请求
            JsonRpcResponse response = proxyService.forwardMcpRequest(key, jsonRpcRequest, proxyInfo);
            
            // 设置会话ID响应头
            if (jsonRpcRequest != null && "initialize".equals(jsonRpcRequest.getMethod())) {
                return ResponseEntity.ok()
                        .header("Mcp-Session-Id", sessionId)
                        .body(response);
            } else {
                return ResponseEntity.ok(response);
            }
        } catch (Exception e) {
            log.error("Error processing MCP request: {}", e.getMessage(), e);
            return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR)
                    .body(JsonRpcResponse.internalError(null, e.getMessage()));
        }
    }
} 